import random
import numpy as np

from beam.util_psll import get_psll_by_phase


class RISAdaptiveGeneticAlgorithm():
    __bit_num = 0  # 比特数
    __beam_num = 0  # 波束数

    def __init__(self, bit_num, beam_num, population_size=50, num_generations=100, num_parents=10, mutation_rate=0.1):
        # 初始化阵列相关
        self.__bit_num = bit_num
        self.__beam_num = beam_num
        # 初始化遗传算法相关
        self.population_size = population_size
        self.num_generations = num_generations
        self.num_parents = num_parents
        self.mutation_rate = mutation_rate
        self.population = None
        self.best_individual = None
        self.best_fitness = None
        self.best_fitness_history = []  # 保存每一代的最佳适应度值
        self.best_individual_history = []  # 保存每一代的最佳适应度值的个体
        # 新增自适应参数
        self.adaptive_params = {
            'pc_min': 0.6,  # 最小交叉概率[7](@ref)
            'pc_max': 0.9,  # 最大交叉概率[7](@ref)
            'pm_base': 0.05,  # 基础变异率[5](@ref)
            'pm_max': 0.2,  # 最大变异率[5](@ref)
            'diversity_threshold': 0.1  # 多样性阈值[6](@ref)
        }

    def __generate_population_part(self, array_1, array_2, max_num):
        part = []
        # 添加before和current作为前两个元素
        part.append(array_1.copy())
        part.append(array_2.copy())
        # 计算 diff 数组
        diff = array_1 != array_2
        # 获取所有 False 的坐标
        false_coords = np.argwhere(diff)
        # 生成剩余的个体
        for _ in range(max_num - 2):
            # 复制 current_array
            modified_array = np.copy(array_2)
            # 随机选择一些差异行
            if np.sum(diff == True) > 0:
                # 随机选择一部分 False 坐标
                num_to_select = np.random.randint(0, len(false_coords) + 1)
                selected_false_coords = false_coords[np.random.choice(len(false_coords), num_to_select, replace=False)]
                # 将 modified_array 中选取坐标替换为 before_array
                for coord in selected_false_coords:
                    row, col = coord
                    modified_array[row, col] = array_1[row, col]
                # print("Modified array:", modified_array.tolist())
            part.append(modified_array)
        return part

    def initialize_population_group(self, phase_init_list):
        part_len = self.population_size // len(phase_init_list)
        self.population = []
        print(f"part_len={part_len}, len: {len(phase_init_list)-1}")
        for i in range(0, len(phase_init_list)):
            phase1 = phase_init_list[i]
            if i < len(phase_init_list) - 1:
                this_len = part_len
                phase2 = phase_init_list[i + 1]
            else:
                this_len = self.population_size - len(self.population)
                phase2 = phase_init_list[0]
            self.population.extend(self.__generate_population_part(phase1, phase2, this_len))

    def initialize_population(self, phase_mix_init):
        """初始化种群"""
        self.population = [phase_mix_init] * self.population_size

    def fitness(self, phase):
        return get_psll_by_phase(phase, self.__bit_num, self.__beam_num)

    def adaptive_selection(self):
        """自适应锦标赛选择"""
        fitness = np.array([self.fitness(ind) for ind in self.population])
        sorted_indices = np.argsort(fitness)
        selected = []

        # 动态锦标赛规模（2-种群大小10%）
        tournament_size = max(2, int(0.1 * self.population_size))
        for _ in range(self.num_parents):
            candidates = random.sample(range(self.population_size), tournament_size)
            winner = min(candidates, key=lambda i: fitness[i])
            selected.append(self.population[winner])

        return selected

    def adaptive_crossover(self, parents):
        """自适应多点交叉"""
        offspring = []
        fitness = np.array([self.fitness(p) for p in parents])
        std_dev = np.std(fitness)

        # 动态交叉率计算
        pc = self.adaptive_params['pc_max'] - \
             (self.adaptive_params['pc_max'] - self.adaptive_params['pc_min']) * \
             (std_dev / (std_dev + 1e-6))

        for _ in range(len(parents) // 2):
            p1, p2 = random.sample(parents, 2)
            if random.random() < pc:
                # 自适应多点交叉
                num_cuts = max(1, int(0.1 * len(p1)))
                cut_points = sorted(random.sample(range(1, len(p1)), num_cuts))
                child1 = np.concatenate([p1[:cut_points[0]],
                                         p2[cut_points[0]:cut_points[1]],
                                         p1[cut_points[1]:]])
                child2 = np.concatenate([p2[:cut_points[0]],
                                         p1[cut_points[0]:cut_points[1]],
                                         p2[cut_points[1]:]])
                offspring.extend([child1, child2])
            else:
                offspring.extend([p1, p2])
        return offspring

    def adaptive_mutation(self, offspring):
        """自适应高斯变异"""
        fitness = np.array([self.fitness(ind) for ind in offspring])
        mean_fitness = np.mean(fitness)

        for i in range(len(offspring)):
            # 基于适应度的变异率调整
            if fitness[i] < mean_fitness:
                current_pm = self.adaptive_params['pm_base'] + \
                             (self.adaptive_params['pm_max'] - self.adaptive_params['pm_base']) * \
                             (1 - fitness[i] / mean_fitness)
            else:
                current_pm = self.adaptive_params['pm_base']

            # 应用自适应高斯变异
            mask = np.random.rand(*offspring[i].shape) < current_pm
            noise = np.random.normal(0, 3, offspring[i].shape)  # 均值0，标准差3
            offspring[i] = np.clip(offspring[i] + mask * noise, -180, 180)
        return offspring

    def elitism_preservation(self, new_population):
        """精英保留策略"""
        combined = self.population + new_population
        fitness = np.array([self.fitness(ind) for ind in combined])
        sorted_indices = np.argsort(fitness)
        return [combined[i] for i in sorted_indices[:self.population_size]]

    def run(self, phase_init_list, logger):
        logger.info(f"GA Parameters: pop={self.population_size}, gens={self.num_generations}, "
                    f"pc=[{self.adaptive_params['pc_min']:.2f},{self.adaptive_params['pc_max']:.2f}], "
                    f"pm=[{self.adaptive_params['pm_base']:.2f},{self.adaptive_params['pm_max']:.2f}]")

        # 初始化种群
        self.initialize_population_group(phase_init_list)
        logger.info(f"len of self.population={len(self.population)}")
        # 输出保存
        fitness_scores = [self.fitness(ind) for ind in self.population]
        this_best_index = np.argmin(fitness_scores)
        self.best_fitness = fitness_scores[this_best_index]
        self.best_individual = self.population[this_best_index]
        self.best_fitness_history = []  # 保存每一代的最佳适应度值
        self.best_individual_history = []  # 保存每一代的最佳适应度值的个体
        # 主循环
        for generation in range(self.num_generations):
            # 自适应选择
            selected_parents = self.adaptive_selection()
            # 自适应交叉
            offspring = self.adaptive_crossover(selected_parents)
            # 自适应变异
            offspring = self.adaptive_mutation(offspring)
            # 精英保留
            self.population = self.elitism_preservation(offspring)
            # 适应度评估
            fitness_scores = [self.fitness(ind) for ind in self.population]
            this_best_index = np.argmin(fitness_scores)
            this_best_fitness = fitness_scores[this_best_index]
            # 更新最优解
            if this_best_fitness < self.best_fitness:
                self.best_fitness = this_best_fitness
                self.best_individual = self.population[this_best_index]

            # 记录历史数据
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)

            # 日志输出
            logger.info(f"Gen {generation:03d}: Best Fitness = {self.best_fitness:.6f} | "
                        f"Diversity = {np.std(fitness_scores):.4f}")

        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history